import torch
import torch.nn as nn
import torch.nn.functional as F
import os
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from model.wholeModel import CombinedModel
from model.dataProcess import SensorDataset,WeightedMSELoss,collate_fn
from tqdm import tqdm


seq_length = 24  # 例如，每个序列包含100个时间步
label_length = 1  # 后20个时间步作为标签
step = 1  # 步长，可以根据需要调整

GAT_hidden_dim = 128
num_sensors = 79
group_size = 10
GAT_output_dim = num_sensors
predict_length = label_length
num_attention_cycles = 5
batch_size = 1

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 检查是否存在保存的模型参数
model_path = 'model\model_GAT_attention_5_24_1.pth'

print("发现保存的模型参数，正在加载...")
model = CombinedModel(seq_length, num_sensors, group_size, GAT_output_dim,GAT_hidden_dim, predict_length, num_attention_cycles)
# model.load_state_dict(torch.load(model_path))
model.load_state_dict(torch.load(model_path, map_location=device))
model.to(device)


# 定义损失函数
criterion = nn.MSELoss(reduction='mean')

# 模型推理
model.eval()

up_threshold = 3#大于这个阈值的loss将被丢弃,用前一时刻的loss代替当前loss
down_threshold = 8#当连续10个时间步的和大于这个阈值，认为异常出现
history = [0.0] * 10#初始化为一个全0数组
def isAttack(batch_x,batch_y):
    outputs = model(batch_x) 
    loss = criterion(outputs, batch_y)
    return loss

def detect_anomaly(batch_x,batch_y):
    global history
    current_loss = isAttack(batch_x,batch_y)
    # 处理当前loss值
    if current_loss > up_threshold:
        current_loss = history[-1]
    
    # 维护最多保留最近10个loss值
    history.append(current_loss)
    history.pop(0)  
    # 检测连续异常条件
    if sum(history) > 8:
        return current_loss,-1#True有异常
    return current_loss,1#False